#include "headers/heap.h"
#include "../tools/headers/common.h"
#include "../api/headers/errors.h"
#include "headers/ordered_array.h"
#include "../memory/headers/kheap.h"
#include "../api/headers/io.h"
#define DEBUG

static void expand(u32int new_size, heap_t *heap);
static u32int contract(u32int new_size, heap_t *heap);
static s32int find_smallest_hole(u32int size, u8int page_align, heap_t *heap);
static s8int header_t_less_than(void*a, void *b);
extern volatile BOOL critical_section;
//Does not expand holes (only allocates pages)
//Errors: ERR_INVALID_ARG - when new size < current size;
//		  ERR_NO_SPACE - when there is no space to expand the heap

static void printHoles(heap_t* heap)
{
	int i;
	printf("Holes:\n");
	header_t* head = 0;
	for(i = 0; i < heap->index.size; i++)
	{
		head = (header_t *)lookup_ordered_array(i, &heap->index);
		printf("[%d] address: 0x%H; size: 0x%H\n", i, head, head->size);
	}
	printf("\n");
}
static void expand(u32int new_size, heap_t *heap)
{

	BOOL backup = critical_section;
	critical_section = true;
    setLastError(NO_ERROR);
#ifdef DEBUG
	printf("\n Expanding heap: 0x%H \n", new_size);
#endif
    //ASSERT(new_size > heap->end_address - heap->start_address);
	if(new_size <= (heap->end_address - heap->start_address))
	{

		setLastError(ERR_INVALID_ARG); //new_size is smaller than current heap size
		return;
	}
    // Get the nearest following page boundary.
    if ((new_size & 0xFFF) != 0)
    {
        new_size &= 0xFFFFF000;
        new_size += 0x1000;
    }

    // Make sure we are not overreaching ourselves.
    //ASSERT(heap->start_address+new_size <= heap->max_address);
	if(heap->start_address + new_size > heap->max_address)
	{

		setLastError(ERR_NO_SPACE);
		return;
	}
#ifdef DEBUG
	printf("Expand2\n");
#endif
	
    // This should always be on a page boundary.
    u32int old_size = heap->end_address - heap->start_address;

    u32int i = old_size;
    while (i < new_size)
    {
        alloc_frame( get_page(heap->start_address + i, 1, heap->directory),
                     (heap->supervisor)?1:0, (heap->readonly)?0:1);
		if(getLastError() != NO_ERROR) //Roll back
		{

		/*	while(i >= old_size)
			{
				free_frame(get_page(heap->start_address + i, 0, heap->directory));
				i -= 0x1000;
			}*/

			return;
		}
        i += 0x1000 /* page size */;
    }
#ifdef DEBUG
	printf("Expand3");
#endif

    heap->end_address = heap->start_address + new_size;
#ifdef DEBUG
	printf("New heap end address: 0x%H\n", heap->end_address);
#endif
	critical_section = backup;
	return;
}

//Contracts heap. Returns new heap size. Contracts only to MIN_HEAP_SIZE
//Errors: ERR_INVALID_ARG - if new size is larger than current size
static u32int contract(u32int new_size, heap_t *heap)
{
	BOOL backup = critical_section;
	critical_section = true;
	setLastError(NO_ERROR);
#ifdef DEBUG
	printf("\n Contracting heap: 0x%H\n", new_size);
#endif
    // Sanity check.
    //ASSERT(new_size < heap->end_address-heap->start_address);
	if(new_size >= (heap->end_address - heap->start_address))
	{
		setLastError(ERR_INVALID_ARG);
		critical_section = backup;
		return heap->end_address - heap->start_address;
	}
    // Get the nearest following page boundary.
    if (new_size & 0xFFF)
    {
        new_size &= 0xFFFFF000;
        new_size += 0x1000;
    }

    // Don't contract too far!
    if (new_size < HEAP_MIN_SIZE)
        new_size = HEAP_MIN_SIZE;

    u32int old_size = heap->end_address - heap->start_address;
    u32int i = old_size - 0x1000;
	printf("i: 0x%H, new_size: 0x%H\n", i, new_size);
    while (new_size <= i)
    {
        free_frame(get_page(heap->start_address + i, 0, heap->directory));
#ifdef DEBUG
		printf("Freeing frame: 0x%H\n", heap->start_address + i);
#endif
        i -= 0x1000;
    }

    heap->end_address = heap->start_address + new_size;
	critical_section = backup;
    return new_size;
}

//Finds the smallest hole (free memory) of the specified size.
//Returns -1 if no hole found.
static s32int find_smallest_hole(u32int size, u8int page_align, heap_t *heap)
{
    // Find the smallest hole that will fit.
    u32int iterator = 0;
    while (iterator < heap->index.size)
    {
        header_t *header = (header_t *)lookup_ordered_array(iterator, &heap->index);
        // If the user has requested the memory be page-aligned
        if (page_align > 0)
        {
            // Page-align the starting point of this header.
            u32int location = (u32int)header;
            s32int offset = 0;
            if (((location + sizeof(header_t)) & 0xFFF) != 0)
                offset = 0x1000 /* page size */  - (location + sizeof(header_t)) % 0x1000;
            s32int hole_size = (s32int)header->size - offset;
            // Can we fit now?
            if (hole_size >= (s32int)size)
                break;
        }
        else if (header->size >= size)
            break;
        iterator++;
    }
    // Why did the loop exit?
    if (iterator == heap->index.size)
	{
        return -1; // We got to the end and didn't find anything.
	}
    else
        return iterator;
}

//Predicate for sorting the array (or list) of all holes
static s8int header_t_less_than(void*a, void *b)
{
    return (((header_t*)a)->size < ((header_t*)b)->size)?1:0;
}

/**
	Creates a heap with specified start and initial end addresses (should be page aligned). Max - maximal heap size. Halts if addresses are not page aligned
	Supervisor, readonly - flags to allocate pages. Assumes that memory is already allocated.

 */
heap_t *create_heap(u32int start, u32int end_addr, u32int max, page_directory_t* directory, u8int supervisor, u8int readonly)
{
	BOOL backup = critical_section;
	critical_section = true;
	setLastError(NO_ERROR);
    heap_t *heap = (heap_t*)kmalloc(sizeof(heap_t));
	
	if(getLastError() != NO_ERROR)
		return 0;
    // All our assumptions are made on startAddress and endAddress being page-aligned.
    ASSERT(start % 0x1000 == 0);
    ASSERT(end_addr % 0x1000 == 0);
    // Initialise the index.
    heap->index = place_ordered_array( (void*)start, HEAP_INDEX_SIZE, &header_t_less_than);
    // Shift the start address forward to resemble where we can start putting data.
    start += sizeof(type_t)*HEAP_INDEX_SIZE;
	
    // Make sure the start address is page-aligned.
    if ((start & 0xFFF) != 0)
    {
        start &= 0xFFFFF000;
        start += 0x1000;
    }
    // Write the start, end and max addresses into the heap structure.
    heap->directory = directory;
    heap->start_address = start;
    heap->end_address = end_addr;
    heap->max_address = max;
    heap->supervisor = supervisor;
    heap->readonly = readonly;
    // We start off with one large hole in the index.
    header_t *hole = (header_t *)start;
    hole->size = end_addr - start;
    hole->magic = HEAP_MAGIC;
    hole->is_hole = 1;
    insert_ordered_array((void*)hole, &heap->index);     
	critical_section = backup;
    return heap;
}

/**
	Allocates specified amount of space, at a specified heap.Expands heap if necessary.Returns a pointer
	to allocated memory.Returns 0 if no memory found.
	\par Errors:
	ERR_NO_MEMORY - when theres is not enough memory in the specified heap
 */
void *alloc(u32int size, u8int page_align, heap_t *heap)
{
	BOOL backup = critical_section;
	critical_section = true;
	setLastError(NO_ERROR);
    // Make sure we take the size of header/footer into account.
    u32int new_size = size + sizeof(header_t) + sizeof(footer_t);
    // Find the smallest hole that will fit.
    s32int iterator = find_smallest_hole(new_size/* + 10*sizeof(page_table_t)*/, page_align, heap);
#ifdef DEBUG
	printf("Allocating 0x%H bytes\n", size);
#endif
    if (iterator == -1) // If we didn't find a suitable hole
    {
        // Save some previous data.
        u32int old_length = heap->end_address - heap->start_address;
        u32int old_end_address = heap->end_address;

        // We need to allocate some more space.
        expand(old_length + new_size/* + 0x1100*/, heap);
#ifdef DEBUG
		printf("Returned after expand\n");
#endif
		if(getLastError() != NO_ERROR)
		{
			critical_section = backup;
			return 0;
		}
		
        u32int new_length = heap->end_address - heap->start_address;

        // Find the endmost header. (Not endmost in size, but in location).
        iterator = 0;
        // Vars to hold the index of, and value of, the endmost header found so far.
        u32int idx = -1; u32int value = 0x0;
        while (iterator < heap->index.size)
        {
            u32int tmp = (u32int)lookup_ordered_array(iterator, &heap->index);
            if (tmp > value)
            {
                value = tmp;
                idx = iterator;
            }
            iterator++;
        }
		header_t* rightHead = (header_t*)value;
		header_t* nextHeader = (header_t*)(value + rightHead->size);
        // If we didn't find ANY headers on right most of the heap, we need to add one.
        if ((idx == -1) || (nextHeader->magic == HEAP_MAGIC))
        {
#ifdef DEBUG
			printf("\n Adding new header at: 0x%H\n", old_end_address);
#endif
            header_t *header = (header_t *)old_end_address;
            header->magic = HEAP_MAGIC;
            header->size = new_length - old_length;
            header->is_hole = 1;
            footer_t *footer = (footer_t *) (old_end_address + header->size - sizeof(footer_t));
            footer->magic = HEAP_MAGIC;
            footer->header = header;
            insert_ordered_array((void*)header, &heap->index);
			if(getLastError() != NO_ERROR)
			{
				return;
			}
			
        }
        else
        {
            // The last header needs adjusting.

            header_t *header = lookup_ordered_array(idx, &heap->index);
#ifdef DEBUG
			printf("\n Adjusting header at: 0x%H\n", header);
#endif
            header->size += new_length - old_length;
            // Rewrite the footer.
            footer_t *footer = (footer_t *) ((u32int)header + header->size - sizeof(footer_t));
            footer->header = header;
            footer->magic = HEAP_MAGIC;
									//UPDATE HOLE IN INDEX LIST (AS SIZE CHANGED)!!
        }
        // We now have enough space. Recurse, and call the function again.
        critical_section = backup;
        return alloc(size, page_align, heap);
    }

    header_t *orig_hole_header = (header_t *)lookup_ordered_array(iterator, &heap->index);
    u32int orig_hole_pos = (u32int)orig_hole_header;
    u32int orig_hole_size = orig_hole_header->size;
    // Here we work out if we should split the hole we found into two parts.
    // Is the original hole size - requested hole size less than the overhead for adding a new hole?
    if (orig_hole_size - new_size < (sizeof(header_t) + sizeof(footer_t)))
    {
        // Then just increase the requested size to the size of the hole we found.
#ifdef DEBUG
        printf("(%d) Expanding new hole size. orig_hole: 0x%H; neew_size: 0x%H\n", __LINE__, orig_hole_size, new_size);
#endif
        size += orig_hole_size - new_size;
        new_size = orig_hole_size;
    }

    // If we need to page-align the data, do it now and make a new hole in front of our block.
    if (page_align && ((orig_hole_pos + sizeof(header_t)) & 0xFFF))
    {
        u32int new_location   = orig_hole_pos + 0x1000 /* page size */ - (orig_hole_pos & 0xFFF) - sizeof(header_t);
		
		if((0x1000 - (orig_hole_pos & 0xFFF) - sizeof(header_t)) > (sizeof(header_t) + 			sizeof(footer_t)))
		{
			header_t *hole_header = (header_t *)orig_hole_pos;
			hole_header->size     = 0x1000 /* page size */ - (orig_hole_pos & 0xFFF) - sizeof(header_t);
			hole_header->magic    = HEAP_MAGIC;
			hole_header->is_hole  = 1;
			footer_t *hole_footer = (footer_t *) ((u32int)new_location - sizeof(footer_t));
			hole_footer->magic    = HEAP_MAGIC;
			hole_footer->header   = hole_header;
			
		
#ifdef DEBUG
			printf("\n Page aligned hole added at: 0x%H\n", new_location);
#endif
			//UPDATE HOLE IN INDEX LIST (AS SIZE CHANGED)!!
		}
		else
		{
			
			footer_t* left_footer = orig_hole_pos - sizeof(footer_t);
			if((left_footer < heap->start_address)?(left_footer->magic != HEAP_MAGIC):false)
			{
				remove_ordered_array(iterator, &heap->index);
				if(getLastError() != NO_ERROR)
				{
					return;
				}
				printf("Warning, creating blank memory area at heap start: 0x%H\n", left_footer);
			}
			else
			{
				header_t* left_header = left_footer->header;
				left_header->size += (new_location - orig_hole_pos);
				footer_t* new_footer = (footer_t*)(new_location - sizeof(footer_t));
				*new_footer = *left_footer;
				left_footer->magic = 0;
			}
			
			//UPDATE HOLE IN INDEX LIST (AS SIZE CHANGED)!!
		}
		
		orig_hole_size        -= (new_location - orig_hole_pos);
		orig_hole_pos         = new_location;
																						
    }
    else
    {
        // Else we don't need this hole any more, delete it from the index.
        remove_ordered_array(iterator, &heap->index);
		if(getLastError() != NO_ERROR)
		{
			return;
		}
    }

    // Overwrite the original header
    header_t *block_header  = (header_t *)orig_hole_pos;
    block_header->magic     = HEAP_MAGIC;
    block_header->is_hole   = 0;
    block_header->size      = new_size;
    // And the footer
    footer_t *block_footer  = (footer_t *) (orig_hole_pos + sizeof(header_t) + size);
    block_footer->magic     = HEAP_MAGIC;
    block_footer->header    = block_header;

    // We may need to write a new hole after the allocated block.
    // We do this only if the new hole would have positive size...
    if ((orig_hole_size - new_size) > (sizeof(header_t) + sizeof(footer_t)))
    {

        header_t *hole_header = (header_t *) (orig_hole_pos + sizeof(header_t) + size + sizeof(footer_t));
        hole_header->magic    = HEAP_MAGIC;
        hole_header->is_hole  = 1;
        hole_header->size     = orig_hole_size - new_size;
        footer_t *hole_footer = (footer_t *) ( (u32int)hole_header + orig_hole_size - new_size - sizeof(footer_t));
        if ((u32int)hole_footer < heap->end_address)
        {
            hole_footer->magic = HEAP_MAGIC;
            hole_footer->header = hole_header;
        }
#ifdef DEBUG
        else
		{
			printf("WARNING! Footer outside heap end_address: %s (%d)\n", __FILE__, __LINE__);
		}
#endif
        // Put the new hole in the index;
        insert_ordered_array((void*)hole_header, &heap->index);
		if(getLastError() != NO_ERROR)
		{
			return;
		}
    }
    critical_section = backup;
    return (void *) ( (u32int)block_header + sizeof(header_t) );
}

void free(void *p, heap_t *heap)
{
	BOOL backup = critical_section;
	critical_section = true;
	setLastError(NO_ERROR);
    // Exit gracefully for null pointers.
    if (p == 0)
        return;

    // Get the header and footer associated with this pointer.
    header_t *header = (header_t*) ( (u32int)p - sizeof(header_t) );
    footer_t *footer = (footer_t*) ( (u32int)header + header->size - sizeof(footer_t) );

    //Checks if memory was allocated
    if((header->magic != HEAP_MAGIC) || (footer->magic != HEAP_MAGIC))
	{
		setLastError(ERR_NOT_ALLOCATED);
		critical_section = backup;
		return;
	}

    // Make us a hole.
    header->is_hole = 1;

    // Do we want to add this header into the 'free holes' index?
    BOOL do_add = true;

    // Unify left
    // If the thing immediately to the left of us is a footer...
    footer_t *test_footer = (footer_t*) ( (u32int)header - sizeof(footer_t) );
    if ((test_footer->magic == HEAP_MAGIC) &&
        (test_footer->header->is_hole == 1))
    {
#ifdef DEBUG
			printf("Unifying left at: 0x%H\n", test_footer);
#endif
        u32int cache_size = header->size; // Cache our current size.
        memset(header, 0, sizeof(header_t));
        header = test_footer->header;     // Rewrite our header with the new one.
        footer->header = header;          // Rewrite our footer to point to the new header.
        header->size += cache_size;       // Change the size.
        do_add = 0;                       // Since this header is already in the index, we don't want to add it again
        memset(test_footer, 0, sizeof(footer_t));
    }

    // Unify right
    // If the thing immediately to the right of us is a header...
    header_t *test_header = (header_t*) ( (u32int)footer + sizeof(footer_t) );
    if ((test_header->magic == HEAP_MAGIC) &&
        test_header->is_hole)
    {
#ifdef DEBUG
			printf("Unifying right at: 0x%H\n", test_header);
#endif
        header->size += test_header->size; // Increase our size.														//UPDATE HOLE IN INDEX LIST (AS SIZE CHANGED)!!
        test_footer = (footer_t*) ((u32int)test_header + // Rewrite it's footer to point to our header.
                                    test_header->size - sizeof(footer_t) );
		memset(footer, 0, sizeof(footer_t));
        footer = test_footer;
        // Find and remove this header from the index.
        u32int iterator = 0;
        while ( (iterator < heap->index.size) &&
                (lookup_ordered_array(iterator, &heap->index) != (void*)test_header))
            iterator++;
		

        // Make sure we actually found the item.
        if(iterator >= heap->index.size)
		{
			setLastError(ERR_HEAP_DAMAGED);
			critical_section = backup;
			return;
		}
		
		memset(test_header, 0, sizeof(header_t));
		footer->header = header;
        remove_ordered_array(iterator, &heap->index);
		if(getLastError() != NO_ERROR)
		{
			return;
		}
    }

    // If the footer location is the end address, we can contract.
    /*if ( ((u32int)footer + sizeof(footer_t)) == heap->end_address)
    {
        u32int old_length = heap->end_address - heap->start_address;
        u32int new_length = contract( (u32int)header - heap->start_address, heap);
		if(getLastError() != NO_ERROR)
		{
			critical_section = backup;
			return;
		}
        // Check how big we will be after resizing.
        if (header->size - (old_length - new_length) > 0)
        {
            // We will still exist, so resize us.
            header->size -= old_length - new_length;
			
            footer = (footer_t*) ( (u32int)header + header->size - sizeof(footer_t) );
            footer->magic = HEAP_MAGIC;
            footer->header = header;
        }
        else
        {
            // We will no longer exist :(. Remove us from the index.
            u32int iterator = 0;
            while ( (iterator < heap->index.size) &&
                    (lookup_ordered_array(iterator, &heap->index) != (void*)test_header) )
                iterator++;
            // If we didn't find ourselves, we have nothing to remove.
            if (iterator < heap->index.size)
			{
                remove_ordered_array(iterator, &heap->index);
                if(getLastError() != NO_ERROR)
				{
					return;
				}
			}
        }
    }*/

    // If required, add us to the index.
    if (do_add == 1)
	{
        insert_ordered_array((void*)header, &heap->index);
		if(getLastError() != NO_ERROR)
		{
			return;
		}
	}
																									//else - UPDATE HOLE IN INDEX LIST (AS SIZE CHANGED)!!
	critical_section = backup;

}

 
